import torch
import torchvision
from torch.optim.lr_scheduler import MultiStepLR
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torchvision import transforms
import model
from model import *
import matplotlib.pyplot as plt

# 设置中文字体（例如：SimHei, Microsoft YaHei）
plt.rcParams['font.sans-serif'] = ['SimHei']  # 黑体
plt.rcParams['axes.unicode_minus'] = False    # 解决负号 '-' 显示为方块的问题

# 创新点1：复合数据增强策略
train_transform = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomRotation(15),
    transforms.ColorJitter(brightness=0.2, contrast=0.2),
    transforms.RandomGrayscale(p=0.1),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2470, 0.2435, 0.2616]),
    transforms.RandomErasing(p=0.5, scale=(0.02, 0.1), ratio=(0.3, 3.3))
])

test_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2470, 0.2435, 0.2616])
])

# 准备数据
train_data = torchvision.datasets.CIFAR10(root='./dataset', train=True, download=True, transform=train_transform)
test_data = torchvision.datasets.CIFAR10(root='./dataset', train=False, download=True, transform=test_transform)

print(f"训练数据长度：{len(train_data)}")
print(f"测试数据长度：{len(test_data)}")

train_loader = DataLoader(train_data, batch_size=128, shuffle=True, num_workers=2)
test_loader = DataLoader(test_data, batch_size=256, shuffle=False, num_workers=2)

# 选择训练设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 实例化网络
my_nn = model.CustomResNet18(True)
# 冻结中间层
freeze_layers(my_nn, num_unfreeze=0)

my_nn = my_nn.to(device)

# 损失函数
loss_fn = nn.CrossEntropyLoss()
loss_fn = loss_fn.to(device)

# 优化器
optimizer = torch.optim.SGD(my_nn.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)

# 学习率调度器
scheduler = MultiStepLR(optimizer, milestones=[40, 70], gamma=0.1)

# 训练轮数
epoch = 100
train_step = 0
test_step = 0

# 记录日志
writer = SummaryWriter("logs_train")

# 初始化最佳准确率
best_accuracy = 0.0

# 初始化数据记录列表
train_loss_list = []
train_acc_list = []
test_loss_list = []
test_acc_list = []

# 训练流程
for i in range(1, epoch + 1):
    print(f"--------------第{i}轮训练开始-----------------")

    # 逐层解冻策略
    if (i - 1) / epoch == 0.4:
        freeze_layers(my_nn, num_unfreeze=2)
        print("解冻layer3-4")
    elif (i - 1) / epoch == 0.7:
        freeze_layers(my_nn, num_unfreeze=4)
        print("解冻layer1-2")

    # 初始化训练指标
    total_train_loss = 0.0
    total_train_correct = 0

    # 训练过程
    for imgs, targets in train_loader:
        imgs = imgs.to(device)
        targets = targets.to(device)

        outputs = my_nn(imgs)
        loss = loss_fn(outputs, targets)
        total_train_loss += loss.item()

        preds = outputs.argmax(1)
        total_train_correct += (preds == targets).sum().item()

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        train_step += 1
        if train_step % 200 == 0:
            writer.add_scalar("train_loss", loss.item(), train_step)

    # 计算训练集平均损失和准确率
    avg_train_loss = total_train_loss / len(train_loader)
    train_acc = total_train_correct / len(train_data)
    train_loss_list.append(avg_train_loss)
    train_acc_list.append(train_acc)

    # 更新学习率
    scheduler.step()
    current_lr = optimizer.param_groups[0]['lr']
    print(f"当前学习率: {current_lr}")

    # 测试过程
    test_total_loss = 0.0
    total_test_correct = 0
    with torch.no_grad():
        for imgs, targets in test_loader:
            imgs = imgs.to(device)
            targets = targets.to(device)

            outputs = my_nn(imgs)
            loss = loss_fn(outputs, targets)
            test_total_loss += loss.item()
            total_test_correct += (outputs.argmax(1) == targets).sum().item()

    # 计算测试集指标
    avg_test_loss = test_total_loss / len(test_loader)
    test_acc = total_test_correct / len(test_data)
    test_loss_list.append(avg_test_loss)
    test_acc_list.append(test_acc)

    # 记录到TensorBoard
    writer.add_scalar("train_loss_epoch", avg_train_loss, i)
    writer.add_scalar("train_accuracy", train_acc, i)
    writer.add_scalar("test_loss", avg_test_loss, i)
    writer.add_scalar("test_accuracy", test_acc, i)

    print(f"第{i}轮训练平均loss：{avg_train_loss:.4f}")
    print(f"第{i}轮测试平均loss：{avg_test_loss:.4f}")
    print(f"第{i}轮训练准确率：{train_acc:.4f}")
    print(f"第{i}轮测试准确率：{test_acc:.4f}")

    # 保存最佳模型
    if test_acc > best_accuracy:
        best_accuracy = test_acc
        torch.save(my_nn.state_dict(), "best_model.pth")
        print(f"模型已更新，当前最佳准确率：{best_accuracy:.4f}")

# 关闭日志记录
writer.close()

# 绘制训练曲线
plt.figure(figsize=(12, 5))

# 绘制准确率曲线
plt.subplot(1, 2, 1)
plt.plot(range(1, epoch+1), train_acc_list, label='训练集准确率')
plt.plot(range(1, epoch+1), test_acc_list, label='测试集准确率')
plt.title('准确率变化曲线')
plt.xlabel('Epoch')
plt.ylabel('准确率')
plt.legend()

# 绘制损失曲线
plt.subplot(1, 2, 2)
plt.plot(range(1, epoch+1), train_loss_list, label='训练集损失')
plt.plot(range(1, epoch+1), test_loss_list, label='测试集损失')
plt.title('损失变化曲线')
plt.xlabel('Epoch')
plt.ylabel('损失')
plt.legend()

plt.tight_layout()
plt.savefig('training_curves.png')
plt.show()